# -*- coding: utf-8 -*-
"""
向xml文件中追加目标框
"""
import sys
import argparse
import os
import cv2
import xml.etree.ElementTree as ET
from xml.dom.minidom import Document


def file_count(directory):
    return sum([len(files) for root,dirs,files in os.walk(directory)])


def is_invalid_big_box(box_bnd,box_rects):
    invalid = False
    in_count = 0
    for box in box_rects:
        if box[0]==box_bnd[0]:
            #判断box是否在box_bnd内部
            if (box_bnd[1]<box[1]) and (box_bnd[2]<box[2]) and (box_bnd[3]>box[3]) and (box_bnd[4]>box[4]):
                in_count +=1
                if in_count > 3:
                    invalid = True
                    break
    return invalid

def calc_rect_area( rect ):
    return (rect[2]-rect[0]+0.001)*(rect[3]-rect[1]+0.001)

def calc_iou(rect1,rect2):
    bd_i= ( max(rect1[0],rect2[0]),max(rect1[1],rect2[1]),\
           min(rect1[2],rect2[2]),min(rect1[3],rect2[3]))
    iw = bd_i[2]-bd_i[0]+0.001
    ih = bd_i[3]-bd_i[1]+0.001
    iou=0
    if( iw>0 and ih>0 ):
        ua = calc_rect_area( rect1 ) + calc_rect_area( rect2 ) - iw*ih
        iou = iw*ih/ua		
    return iou

def add_object(doc,annotation,obj_name,xmin_v,ymin_v,xmax_v,ymax_v,occl):    
    obj= doc.createElement("object")   
    objname = doc.createElement("name")
    objname.appendChild(doc.createTextNode(obj_name))
    pose = doc.createElement("pose")
    pose.appendChild(doc.createTextNode("front"))
    truncated = doc.createElement("truncated")
    truncated.appendChild(doc.createTextNode(str(occl)))
    difficult = doc.createElement('difficult')
    difficult.appendChild(doc.createTextNode("0"))
    obj.appendChild(objname)
    obj.appendChild(pose)
    obj.appendChild(truncated)
    obj.appendChild(difficult)
    
    bndbox = doc.createElement("bndbox")
    xmin = doc.createElement("xmin")
    ymin = doc.createElement("ymin")
    xmax = doc.createElement("xmax")
    ymax = doc.createElement("ymax")
    xmin.appendChild(doc.createTextNode(str(xmin_v)))
    ymin.appendChild(doc.createTextNode(str(ymin_v)))
    xmax.appendChild(doc.createTextNode(str(xmax_v)))
    ymax.appendChild(doc.createTextNode(str(ymax_v)))
    bndbox.appendChild(xmin)
    bndbox.appendChild(ymin)
    bndbox.appendChild(xmax)
    bndbox.appendChild(ymax)
    obj.appendChild(bndbox)
    
    annotation.appendChild(obj)
 

def generate_truth_dict(append_truth_root, append_classname):
    ground_truth_dict = {}
    count = 0
    for root, dirs, files in os.walk(append_truth_root):
        for truth_file in files:
            tfile = os.path.join(append_truth_root, truth_file)
            img_name = truth_file.split('.')[0]
            ground_truth_dict[img_name] = []
            tf = open(tfile)
            for line in tf.readlines():
                info = line.split(',')
                info = list(map(int,info))
                ground_truth_dict[img_name].append(info)
                    
            count += 1
            if count % 500 == 0:
                print("generate truth: %d" % count)
    return ground_truth_dict
            
    
def append_object(append_truth_root, voc_root, append_classname, add_classname, img_suffix):
    img_path = os.path.join(voc_root, 'JPEGImages/')
    annotations_path = os.path.join(voc_root, 'Annotations/')
    if not os.path.exists(annotations_path):
        os.makedirs(annotations_path)
    append_annotations_path = os.path.join(voc_root, 'append_Annotations/')
    if not os.path.exists(append_annotations_path):
        os.makedirs(append_annotations_path)
    
    img_num = file_count(img_path)
    xml_num = file_count(annotations_path)
    append_truth_count = file_count(append_truth_root)

    ground_truth_dict = generate_truth_dict(append_truth_root, append_classname)
    
    complete_count = 0
    for img_name in ground_truth_dict:
        img_file = None
        for suffix in img_suffix:
            if os.path.exists(os.path.join(img_path, img_name+suffix)):
                img_file = os.path.join(img_path, img_name+suffix)
                break
        if img_file is None:
            print("image %s not exists!"%img_name)
            continue
        doc = Document()
        img = cv2.imread(img_file)
        n_rows,n_cols,n_channels = img.shape
        
        annotation = doc.createElement('annotation')
        doc.appendChild(annotation)
        
        folder = doc.createElement('folder')
        folder.appendChild(doc.createTextNode("AIA AUTO"))
        annotation.appendChild(folder)
        
        filename = doc.createElement('filename')
        filename.appendChild(doc.createTextNode(img_file.split('/')[-1]))
        annotation.appendChild(filename)
        
        source = doc.createElement('source')
        database = doc.createElement('database')
        database.appendChild(doc.createTextNode('The AUTO Database'))
        anno = doc.createElement("annotation")
        anno.appendChild(doc.createTextNode("HIKDATA_AIA AUTO"))
        image = doc.createElement("image")
        image.appendChild(doc.createTextNode("flickr"))
        source.appendChild(database)
        source.appendChild(anno)
        source.appendChild(image)
        annotation.appendChild(source)
        
        sizeimage = doc.createElement('size')
        imagewidth = doc.createElement('width')
        imagewidth.appendChild(doc.createTextNode(str(n_cols)))
        imageheight = doc.createElement('height')
        imageheight.appendChild(doc.createTextNode(str(n_rows)))
        imagedepth = doc.createElement("depth")
        imagedepth.appendChild(doc.createTextNode(str(n_channels)))
        sizeimage.appendChild(imagewidth)
        sizeimage.appendChild(imageheight)
        sizeimage.appendChild(imagedepth)
        annotation.appendChild(sizeimage)
        
        exist_bndx = {}
        
        '''原本目标框'''
        if os.path.exists(os.path.join(annotations_path, img_name+'.xml')):
            
            tree = ET.parse(os.path.join(annotations_path,img_name+'.xml'))
            root = tree.getroot()
            for obj in root.iter('object'):
                obj_name = obj.find('name').text
                occl = int(float(obj.find('truncated').text))
                xml_box = obj.find('bndbox')
                xmin_v = int(xml_box.find('xmin').text)
                ymin_v = int(xml_box.find('ymin').text)
                xmax_v= int(xml_box.find('xmax').text)
                ymax_v= int(xml_box.find('ymax').text)
                if xmax_v > xmin_v and ymax_v >ymin_v:
                    add_object(doc,annotation,obj_name,xmin_v,ymin_v,xmax_v,ymax_v,occl)
                else:
                    print(os.path.join(annotations_path,img_name+'.xml'))
                    print("xmax_v<=xmin_v or ymax_v<=ymin_v")
                    continue
                #检查类别是否存在
                if obj_name not in exist_bndx:
                    print(os.path.join(annotations_path,img_name+'.xml'))
                    print(obj_name," not in exist class,please check error lable or intentional selection?")
                    continue
                if add:
                    exist_bndx[obj_name].append([xmin_v,ymin_v,xmax_v,ymax_v])
                
        
        '''追加的目标框'''
        for box in ground_truth_dict[img_name]:
            if box[0] in append_classname:
                obj_name = append_classname[box[0]]
                xmin_v = box[1]
                ymin_v = box[2]
                xmax_v = box[3]
                ymax_v = box[4]
                '''无效大框'''
                if ((xmax_v-xmin_v)>n_cols/3) or ((ymax_v-ymin_v)>n_rows/3):
                    if is_invalid_big_box(box,ground_truth_dict[img_name]):
                        continue
                occl=0
                add_object(doc,annotation,obj_name,xmin_v,ymin_v,xmax_v,ymax_v,occl)
            elif (box[0] in add_classname) and add:
                obj_name = add_classname[box[0]]
                box_bd = box[1:]
                sure_add=True
                for rect in exist_bndx[obj_name]:
                    iou = calc_iou(rect,box_bd)
                    if iou>0.1:
                        sure_add = False
                        break
                if sure_add:
                    add_object(doc,annotation,obj_name,box_bd[0],box_bd[1],box_bd[2],box_bd[3],0)
                                          
                              
        xml_name = img_name+'.xml'
        xml_file = open(os.path.join(append_annotations_path, xml_name), 'w')
        xml_file.write(doc.toprettyxml(indent=' '*4))
        complete_count += 1
        if complete_count%1000==0:
            print('[%d | %d] %d%%'%(append_truth_count,complete_count,complete_count*100/append_truth_count))
    print('[%d | %d] %d%%'%(append_truth_count,complete_count,complete_count*100/append_truth_count))
    
def main():
    append_classname = {0: 'vehicle'}  # 追加的类别,前面序号为检测算法对应的类别序号
    exist_classname={1:'person'} #xml中已经存在的目标类别
    add_classname={1:'person'} #添加漏标框，标注较差的
    add = True #是否添加漏标框
    append_truth_root='/mnt/D/hiksense/zxl/yolov3/darknet-master/intermediate/'
    voc_root='/mnt/D/hiksense/dataset/PASCAL_VOC/CUHK_Occlusion/voc'
    img_suffix=['.jpg','.png','.bmp'] #可能的图像后缀
    append_object(append_truth_root,voc_root,append_classname,exist_classname,add_classname,add,img_suffix)


if __name__ == '__main__':
    try:
        ROOT_DIR = sys.argv[1]
        EXTRA_OBJ_DIR = float(sys.argv[2])
    except:
        PARSER = argparse.ArgumentParser(description="arguments")
        PARSER.add_argument("--root_dir", type=str,
                            default="/home/chujie/PycharmProjects/data_tool/data/voc2012",
                            help="VOC dataset root directory")
        PARSER.add_argument("--extra_obj_dir", type=str, default="data/voc2012/result",
                            help="extra objects to be added")
        ARGS = PARSER.parse_args()
        ROOT_DIR = ARGS.root_dir
        EXTRA_OBJ_DIR = ARGS.extra_obj_dir
    # 追加的类别,前面序号为检测算法对应的类别序号
    append_classname = {0: 'vehicle'}
    add_classname = {1: 'person'}  # 添加漏标框，标注较差的
    img_suffix = ['.jpg', '.png', '.bmp']
    append_object(EXTRA_OBJ_DIR, ROOT_DIR, append_classname, add_classname, img_suffix)
